import * as tf from "@tensorflow/tfjs";
import * as tfvis from "@tensorflow/tfjs-vis";

(async () => {
  let weights = [50, 60, 65, 78, 65, 99, 88, 75, 65, 46, 96, 56];
  let height = [150, 160, 172, 181, 160, 192, 165, 168, 156, 180, 176,165];

  tfvis.render.scatterplot(
    {
      name: "信息",
    },
    {
      values: height.map((val, index) => ({ x: val, y: weights[index] })),
    },
    {
      xAxisDomain: [140, 200],
    }
  );

  let model = tf.sequential();
  model.add(
    tf.layers.dense({
      units: 1,
      inputShape: [1],
    })
  );

  model.compile({
    loss: tf.losses.meanSquaredError, // 均方误差
    optimizer: tf.train.sgd(0.01), //
  });
  // 先剪掉最小值，之后除以最大值 - 最小值
  let inputs = tf.tensor(height).sub(150).div(31);
  let labels = tf.tensor(weights).sub(46).div(53);

  await model.fit(inputs, labels, {
    // 学习样本数量
    batchSize: 5,
    // 迭代多少次训练次数
    epochs: 200,
    // 调用
    callbacks: tfvis.show.fitCallbacks(
      { name: "训练" },
      // 度量单位可视化单位
      ["loss"]
    ),
  });
  // 要预测的值也要归一化
  let aa = model.predict(tf.tensor([180]).sub(150).div(31));
  // 输出也要归一化
  console.log( aa.mul(53).add(46).dataSync());
})();
